import requests
import pandas as pd
from datetime import datetime
import json
import logging
import sys
from pyhive import hive
import traceback
import concurrent.futures
from concurrent.futures import ThreadPoolExecutor


class WeatherDataCollector:
    def __init__(self, api_key, locations):
        self.api_key = api_key
        self.locations = locations
        self.base_url = "http://api.weatherapi.com/v1"
        self.hive_config = {
            'host': '192.168.88.161',
            'port': 10000,
            'database': 'weather_db',
            'auth': 'NOSASL'
        }

    def _get_hive_connection(self):
        """获取Hive连接"""
        try:
            print("连接Hive数据库...")
            conn = hive.Connection(**self.hive_config)
            print("Hive连接成功！")
            return conn
        except Exception as e:
            print(f"Hive连接失败: {str(e)}")
            print(traceback.format_exc())
            return None

    def _insert_to_hive(self, df):
        """将数据批量插入Hive表"""
        conn = None
        try:
            conn = self._get_hive_connection()
            if not conn:
                raise Exception("无法连接到Hive数据库")

            cursor = conn.cursor()

            # 设置Hive动态分区模式
            print("\n配置Hive参数...")
            cursor.execute("SET hive.exec.dynamic.partition.mode=nonstrict")
            cursor.execute("SET hive.exec.dynamic.partition=true")
            cursor.execute("SET hive.exec.max.dynamic.partitions=10000")
            cursor.execute("SET hive.exec.max.dynamic.partitions.pernode=10000")

            # 准备批量插入的数据
            print("\n准备批量插入数据...")
            values_list = []

            for _, row in df.iterrows():
                # 解析时间获取年月
                dt = datetime.strptime(row['timestamp'], '%Y-%m-%d %H:%M')
                year = dt.year
                month = dt.month

                # 准备数据
                values = (
                    row['timestamp'],  # date_time
                    row['city'],  # location
                    float(row['temperature']),
                    float(row['humidity']),
                    float(row['pressure']),
                    float(row['wind_speed']),
                    row['wind_direction'],
                    float(row['precipitation']),
                    row['weather_condition'],
                    year,
                    month
                )
                values_list.append(values)

            # 构建批量插入语句
            insert_sql = """
            INSERT INTO TABLE weather_data
            PARTITION (year, month)
            VALUES
            """

            # 每1000条数据一批
            batch_size = 1000
            for i in range(0, len(values_list), batch_size):
                batch = values_list[i:i + batch_size]
                values_str = ",".join([
                    f"('{v[0]}', '{v[1]}', {v[2]}, {v[3]}, {v[4]}, {v[5]}, '{v[6]}', {v[7]}, '{v[8]}', {v[9]}, {v[10]})"
                    for v in batch
                ])

                print(f"\n插入第 {i // batch_size + 1} 批数据 ({len(batch)} 条)...")
                cursor.execute(insert_sql + values_str)

            print("数据插入完成")

            # 更新统计表
            print("\n更新每日统计数据...")
            daily_stats_query = """
            INSERT OVERWRITE TABLE daily_weather_stats
            SELECT 
                location,
                DATE(date_time) as dt,
                MAX(temperature) as max_temp,
                MIN(temperature) as min_temp,
                AVG(temperature) as avg_temp,
                AVG(humidity) as avg_humidity,
                SUM(precipitation) as total_precipitation,
                year,
                month
            FROM weather_data
            GROUP BY location, DATE(date_time), year, month
            """
            cursor.execute(daily_stats_query)

            print("\n更新月度统计数据...")
            monthly_stats_query = """
            INSERT OVERWRITE TABLE monthly_weather_stats
            SELECT 
                location,
                year,
                month,
                MAX(temperature) as max_temp,
                MIN(temperature) as min_temp,
                AVG(temperature) as avg_temp,
                AVG(humidity) as avg_humidity,
                SUM(precipitation) as total_precipitation
            FROM weather_data
            GROUP BY location, year, month
            """
            cursor.execute(monthly_stats_query)

            print("所有数据更新完成！")
            return True

        except Exception as e:
            print(f"插入数据到Hive时出错: {str(e)}")
            print(traceback.format_exc())
            return False

        finally:
            if conn:
                try:
                    conn.close()
                    print("数据库连接已关闭")
                except:
                    pass

    def collect_city_data(self, city):
        """收集单个城市的天气数据"""
        try:
            print(f"\n正在收集 {city} 的天气数据...")

            # 收集当前数据
            current_url = f"{self.base_url}/current.json"
            current_params = {
                'key': self.api_key,
                'q': city,
                'aqi': 'yes'
            }

            # 收集历史数据（过去24小时）
            history_url = f"{self.base_url}/history.json"
            today = datetime.now().strftime('%Y-%m-%d')
            history_params = {
                'key': self.api_key,
                'q': city,
                'dt': today
            }

            # 获取数据
            current_response = requests.get(current_url, params=current_params)
            history_response = requests.get(history_url, params=history_params)

            current_response.raise_for_status()
            history_response.raise_for_status()

            current_data = current_response.json()
            history_data = history_response.json()

            # 处理数据
            weather_records = []
            now = datetime.now()

            # 处理当前数据
            current = current_data['current']
            location = current_data['location']

            # 使用精确的时间戳
            current_time = datetime.strptime(location['localtime'], '%Y-%m-%d %H:%M')
            current_record = {
                'timestamp': current_time.strftime('%Y-%m-%d %H:%M'),
                'city': location['name'],
                'temperature': current['temp_c'],
                'humidity': current['humidity'],
                'pressure': current['pressure_mb'],
                'wind_speed': current['wind_kph'],
                'wind_direction': current['wind_dir'],
                'precipitation': current['precip_mm'],
                'weather_condition': current['condition']['text']
            }
            weather_records.append(current_record)

            # 处理历史数据
            hour_data = history_data['forecast']['forecastday'][0]['hour']

            # 选择最近24小时内的9个均匀分布的时间点
            total_hours = len(hour_data)
            step = total_hours // 9  # 将24小时分成9段

            for i in range(0, total_hours, step):
                if len(weather_records) >= 10:  # 每个城市最多10条记录
                    break

                hour = hour_data[i]
                hour_time = datetime.strptime(hour['time'], '%Y-%m-%d %H:%M')

                if hour_time < now:  # 只取过去的时间点
                    history_record = {
                        'timestamp': hour_time.strftime('%Y-%m-%d %H:%M'),
                        'city': location['name'],
                        'temperature': hour['temp_c'],
                        'humidity': hour['humidity'],
                        'pressure': hour['pressure_mb'],
                        'wind_speed': hour['wind_kph'],
                        'wind_direction': hour['wind_dir'],
                        'precipitation': hour['precip_mm'],
                        'weather_condition': hour['condition']['text']
                    }
                    weather_records.append(history_record)

            # 按时间排序
            weather_records.sort(key=lambda x: datetime.strptime(x['timestamp'], '%Y-%m-%d %H:%M'))

            print(f"已收集 {city} 的天气数据 ({len(weather_records)} 条记录)")
            return weather_records

        except Exception as e:
            print(f"收集 {city} 数据时出错: {str(e)}")
            return []

    def collect_weather_data(self):
        """并行收集所有城市的天气数据"""
        try:
            all_weather_data = []

            # 使用线程池并行收集数据
            with ThreadPoolExecutor(max_workers=min(10, len(self.locations))) as executor:
                # 提交所有任务
                future_to_city = {
                    executor.submit(self.collect_city_data, city): city
                    for city in self.locations
                }

                # 获取结果
                for future in concurrent.futures.as_completed(future_to_city):
                    city = future_to_city[future]
                    try:
                        city_data = future.result()
                        all_weather_data.extend(city_data)
                    except Exception as e:
                        print(f"处理 {city} 数据时出错: {str(e)}")

            if not all_weather_data:
                print("警告：没有收集到任何数据！")
                return None

            # 转换为DataFrame
            df = pd.DataFrame(all_weather_data)
            print(f"\n成功收集了 {len(all_weather_data)} 条数据")

            # 保存到Hive
            if self._insert_to_hive(df):
                print("所有数据处理完成")
            else:
                print("警告：数据保存到Hive失败")

            return df

        except Exception as e:
            print(f"收集天气数据时出错: {str(e)}")
            print(traceback.format_exc())
            return None


if __name__ == "__main__":
    try:
        # 只保留5个主要城市
        cities = [
            'Beijing', 'Shanghai', 'Guangzhou', 'Shenzhen', 'Hangzhou'
        ]

        # 从配置文件读取API密钥
        with open('config.json', 'r') as f:
            config = json.load(f)

        api_key = config['WeatherAPI']['api_key']

        print(f"准备收集以下城市的天气数据: {', '.join(cities)}")

        collector = WeatherDataCollector(api_key, cities)
        data = collector.collect_weather_data()

        if data is not None:
            print("\n数据收集和导入完成")

    except Exception as e:
        print(f"程序运行出错: {str(e)}")
        print(traceback.format_exc())
        sys.exit(1)